import os.path
import time,h5py
from threading import Thread
from PyQt5.QtWidgets import QWidget,QFileDialog,QAbstractItemView,QToolTip
from PyQt5.QtGui import QIntValidator,QDoubleValidator,QStandardItem,QStandardItemModel,QCursor
from PyQt5.QtCore import Qt,pyqtSignal,QTimer
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg,NavigationToolbar2QT
from matplotlib.figure import Figure
from gui.windows.profileWin import Ui_profileWindow
from config.method import *
from config.constant import *
from data.h5_wr import get_profile_data

class ProfileWindow(QWidget, Ui_profileWindow):

    signal_update_process_text = pyqtSignal(str,bool)
    signal_profile_save = pyqtSignal(int)
    signal_update_profile_bar=pyqtSignal(float)
    signal_timer_result_show=pyqtSignal(bool)

    def __init__(self,startWindow):
        super(ProfileWindow, self).__init__()
        self.setupUi(self)
        self.setFixedSize(1381, 800)
        self.startWindow=startWindow
        self.param_dict=startWindow.param_dict
        self.initial_property()
        self.initial_widget()
        self.connect_slot()

    def initial_property(self):
        self.text_string = ''
        self.profile_view_bytes = []
        self.continue_train = False
        self.continue_train_path = None
        self.figure = Figure()
        self.figure.subplots_adjust(top=0.85, wspace=0.53)
        self.canvas = FigureCanvasQTAgg(self.figure)
        self.toolbar = NavigationToolbar2QT(self.canvas, self)
        self.vL_result_show.addWidget(self.canvas)
        self.hL_toolbar.addWidget(self.toolbar)
        attacker_class = type(self.param_dict['attack_method'] + "_class", (object,), {
            'profile': func_profile[self.param_dict['attack_method']],
            'param_dict': self.param_dict,
            'traces_profile': None,
            'midvalue_profile': None
        })
        self.attacker = attacker_class()
        with open("resource/style.qss", encoding='UTF-8') as f:
            self.setStyleSheet(f.read())

    def initial_widget(self):
        self.attack_method.addItems(ATTACK_METHOD_PROFILE)
        self.network.addItems(NETWORK)
        self.epochs.setText(str(self.param_dict['epochs']))
        self.batch_size.setText(str(self.param_dict['batch_size']))
        self.learning_rate.setText(str(self.param_dict['learning_rate']))
        self.activation.addItems(['relu','selu','tanh','elu'])
        self.activation.setCurrentText(self.param_dict['activation'])
        self.validation_split.setValue(self.param_dict["validation_split"])
        self.profile_bar.setMaximum(100)
        self.profile_bar.setValue(0)
        self.profile_bar.hide()
        self.stop_train.setEnabled(False)
        self.profile_path_view.horizontalHeader().setFixedHeight(30)
        self.profile_path_view.resizeColumnsToContents()
        self.profile_path_view.setSelectionMode(QAbstractItemView.SingleSelection)
        self.profile_path_model = QStandardItemModel(0, 2)
        self.profile_path_model.setHorizontalHeaderLabels(['目标字节', '建模数据集路径'])
        self.profile_path_view.setModel(self.profile_path_model)
        self.timer_result_show=QTimer(self)
        self.set_validator()

    def connect_slot(self):
        self.attack_method.currentTextChanged.connect(self.set_attack_method)
        self.add_item.clicked.connect(self.add_profile_item)
        self.remove_item.clicked.connect(self.remove_profile_item)
        self.start_profile.clicked.connect(self.start_profile_process)
        self.clear_related_info.clicked.connect(lambda :self.update_text('#'))
        self.profile_path_view.doubleClicked.connect(self.modify_profile_path)
        self.profile_path_view.clicked.connect(self.get_item_tip)
        self.profile_path_model.itemChanged.connect(self.item_check)
        self.train_new_network.clicked.connect(self.set_train_new_network)
        self.exist_model_path_browse.clicked.connect(self.browse_exist_model_path)
        self.stop_train.clicked.connect(self.set_stop_train)
        self.show_valid_result.clicked.connect(self.update_result_show)
        self.timer_result_show.timeout.connect(self.timer_update_result_show)
        self.view_dataset_info.clicked.connect(self.viewing_dataset_info)
        self.related_info.textChanged.connect(lambda: self.related_info.moveCursor(11))
        self.signal_update_process_text.connect(self.update_process_text)
        self.signal_profile_save.connect(self.save_profile)
        self.signal_update_profile_bar.connect(self.update_profile_bar)
        self.signal_timer_result_show.connect(self.ctrl_timer_result_show)

    def set_validator(self):
        int_validator = QIntValidator(10, 999, self)
        double_validator=QDoubleValidator(0,1,6,self)
        double_validator.setNotation(QDoubleValidator.StandardNotation)
        self.epochs.setValidator(int_validator)
        self.batch_size.setValidator(int_validator)
        self.learning_rate.setValidator(double_validator)

    def fresh_param(self):
        self.param_dict["network"]=self.network.currentText()
        self.param_dict["epochs"]=int(self.epochs.text())
        self.param_dict["batch_size"]=int(self.batch_size.text())
        self.param_dict["learning_rate"]=float(self.learning_rate.text())
        self.param_dict["activation"]=self.activation.currentText()
        self.param_dict["validation_split"] = float(self.validation_split.value())

    def set_attack_method(self,method):
        if self.param_dict["attack_method"]!=method:
            self.reset_attacker(method)
        self.param_dict["attack_method"] = method
        switch_hide1=method=='TA'
        switch_hide2=switch_hide1 or self.train_new_network.isChecked() is False
        switch_hide3=switch_hide1 or self.train_new_network.isChecked() is True
        self.train_new_network.setHidden(switch_hide1)
        self.validation_split_L.setHidden(switch_hide1)
        self.validation_split.setHidden(switch_hide1)
        self.autosave_hyper.setHidden(switch_hide1)
        self.autosave_figure.setHidden(switch_hide1)
        self.network_L.setHidden(switch_hide2)
        self.network.setHidden(switch_hide2)
        self.epochs_L.setHidden(switch_hide2)
        self.epochs.setHidden(switch_hide2)
        self.batch_size_L.setHidden(switch_hide2)
        self.batch_size.setHidden(switch_hide2)
        self.learning_rate_L.setHidden(switch_hide2)
        self.learning_rate.setHidden(switch_hide2)
        self.activation_L.setHidden(switch_hide2)
        self.activation.setHidden(switch_hide2)
        self.exist_model_path_L.setHidden(switch_hide3)
        self.exist_model_path.setHidden(switch_hide3)
        self.exist_model_path_browse.setHidden(switch_hide3)
        if switch_hide1:
            self.autosave.setText("自动保存模板") #模板攻击
        else:
            self.autosave.setText("自动保存模型")  #深度学习

    def reset_attacker(self,method=None):
        attacker_class = type(method + "_class", (object,), {
            'profile': func_profile[method],
            'param_dict': self.param_dict,
            'traces_profile': self.attacker.traces_profile,
            'midvalue_profile': self.attacker.midvalue_profile
        })
        self.attacker = attacker_class()

    def save_profile(self,byte):
        ctime = time.localtime()
        dir = self.param_dict['model_file_dir'] + f"{ctime.tm_year}-{ctime.tm_mon}-{ctime.tm_mday}/"
        if os.path.exists(dir) is False:
            os.makedirs(dir)
        hms = f'{ctime.tm_hour:02d}{ctime.tm_min:02d}{ctime.tm_sec:02d}'
        try:
            if self.attack_method.currentText()=='TA':
                file_name = f"TA_{self.param_dict['target_algorithm']}_{self.param_dict['leakage_model']}_byte{byte}_{hms}.h5"
                if self.autosave.isChecked():
                    fname=dir+file_name
                else:
                    fname,_ = QFileDialog.getSaveFileName(self, "Save file", dir+file_name, "*.h5")
                if fname:
                    h5_file = h5py.File(fname, "w")
                    h5_file.create_dataset(name="template", data=self.attacker.template)
                    h5_file.flush()
                    h5_file.close()
                    self.update_text(f"目标字节{byte}对应模板已保存到{fname}",False)
            if self.attack_method.currentText()!='TA':
                file_name = f"{self.param_dict['network']}_{self.param_dict['target_algorithm']}_{self.param_dict['leakage_model']}_byte{byte}_{hms}.h5"
                if self.autosave.isChecked():
                    fname = dir + file_name
                else:
                    fname,_ = QFileDialog.getSaveFileName(self, "Save file", dir + file_name, "*.h5")
                if fname:
                    self.attacker.model.save(filepath=fname, overwrite=True, include_optimizer=True)
                    self.update_text(f"目标字节{byte}对应模型已保存到{fname}", False)
        except Exception as e:
            self.update_text(repr(e), False)
            self.update_text(f"目标字节{byte}对应模型保存失败！", False)
        self.continue_profile=True

    def ctrl_timer_result_show(self,start):
        if start:
            self.timer_result_show.start(100)
        else:
            self.profile_bar.setValue(100)
            self.timer_result_show.stop()
            ctime=time.localtime()
            dir = self.param_dict['model_file_dir'] + f"{ctime.tm_year}-{ctime.tm_mon}-{ctime.tm_mday}/"
            if os.path.exists(dir) is False:
                os.makedirs(dir)
            hms = f'{ctime.tm_hour:02d}{ctime.tm_min:02d}{ctime.tm_sec:02d}'
            if self.autosave_figure.isChecked():
                file_name = f"{self.param_dict['network']}_{self.param_dict['target_algorithm']}_{self.param_dict['leakage_model']}_byte{self.cur_byte}_{hms}.{self.param_dict['figure_format']}"
                fig_save_path = dir + file_name
                self.figure.savefig(fig_save_path, dpi=self.param_dict["figure_dpi"])
                self.update_text(f"训练结果图片已保存至{fig_save_path}", False)
            if self.autosave_hyper.isChecked():
                file_name=f"{self.param_dict['network']}_{self.param_dict['target_algorithm']}_{self.param_dict['leakage_model']}_byte{self.cur_byte}_{hms}.txt"
                with open(dir+file_name,"w") as f:
                    data=f"train_time: {ctime.tm_year}-{ctime.tm_mon}-{ctime.tm_mday} {ctime.tm_hour:02d}:{ctime.tm_min:02d}:{ctime.tm_sec:02d}\n"
                    data+=f"target algorithm: {self.param_dict['target_algorithm']}\n"
                    data+=f"leakage model: {self.param_dict['leakage_model']}\n"
                    data+=f"target byte: {self.cur_byte}\n"
                    data+=f"network: {self.param_dict['network']}\n"
                    data+=f"epochs: {len(self.attacker.accuracy)}\n"
                    data+=f"batch size: {self.param_dict['batch_size']}\n"
                    data+=f"learning rate: {self.param_dict['learning_rate']}\n"
                    data+=f"activation: {self.param_dict['activation']}\n"
                    data+=f"validation split: {self.param_dict['validation_split']}\n"
                    data+=f"accuracy: {self.attacker.accuracy[-1]}\n"
                    data+=f"loss: {self.attacker.loss[-1]}\n"
                    data += f"accuracy list: {self.attacker.accuracy}\n"
                    data += f"loss list: {self.attacker.loss}\n"
                    if self.attacker.val_loss is not None and len(self.attacker.val_loss)>0:
                        data+=f"val_accuracy: {self.attacker.val_accuracy[-1]}\n"
                        data+=f"val_loss: {self.attacker.val_loss[-1]}\n"
                        data += f"val_accuracy list: {self.attacker.val_accuracy}\n"
                        data += f"val_loss list: {self.attacker.val_loss}\n"
                    f.write(data)
                    self.update_text(f"训练参数已保存至{dir+file_name}", False)


    def start_profile_th(self):
        tip="建模完毕."
        for i in range(len(self.profile_path_list)):
            try:
                self.continue_profile=False
                self.cur_byte=self.profile_bytes_list[i]
                traces, midvalue,description=get_profile_data(self.profile_path_list[i])
                self.attacker.traces_profile=traces
                self.attacker.midvalue_profile=midvalue
                description=description.replace(',','\n').replace('{','{\n').replace('}','\n}')
                self.signal_update_process_text.emit(f"建模数据集信息：\n{description}", False)
                if self.param_dict['attack_method']=="TA":
                    ta_gen=self.attacker.profile()
                    while True:
                        bar_value=next(ta_gen)
                        self.signal_update_profile_bar.emit(bar_value)
                else:
                    self.attacker.loss, self.attacker.accuracy = [], []
                    self.attacker.val_loss, self.attacker.val_accuracy = [], []
                    self.attacker.stop_train = False
                    self.result_list_len=0
                    self.signal_timer_result_show.emit(True)
                    self.attacker.profile()
                    self.signal_timer_result_show.emit(False)
                    self.attacker.stop_train = False
                    self.signal_profile_save.emit(self.profile_bytes_list[i])
                    self.signal_update_process_text.emit(f"目标字节{self.cur_byte}对应模型建模完毕.", False)
                while self.continue_profile is False:
                    time.sleep(0.1)
            except StopIteration:
                self.signal_update_process_text.emit(f"目标字节{self.cur_byte}对应模板生成完毕.", False)
                self.signal_profile_save.emit(self.cur_byte)
            except Exception as e:
                self.signal_update_process_text.emit(repr(e),False)
                self.signal_update_process_text.emit(f"目标字节{self.cur_byte}建模失败！", False)
                tip="建模失败！"
        self.signal_update_process_text.emit(tip, True)

    def start_profile_process(self):
        if self.param_dict['attack_method']!='TA':
            self.fresh_param()
            self.attacker.continue_train=self.continue_train
            self.attacker.continue_train_path=self.continue_train_path
            self.attacker.stop_train=False
            self.stop_train.setEnabled(True)
        self.profile_bytes_list=[]
        self.profile_path_list=[]
        for i in range(self.profile_path_model.rowCount()):
            path=self.profile_path_model.item(i,1).text()
            if path!='未设置' and os.path.exists(path):
                self.profile_bytes_list.append(int(self.profile_path_model.item(i,0).text()))
                self.profile_path_list.append(path)
        if len(self.profile_path_list)==0:
            self.update_text("请设置正确的地址！",False)
            return
        self.update_text("开始建模...", False)
        self.start_profile.setEnabled(False)
        self.attack_method.setEnabled(False)
        self.profile_bar.show()
        th_profile=Thread(target=self.start_profile_th)
        th_profile.start()

    def set_stop_train(self):
        self.attacker.stop_train=True

    def item_check(self,item):
        if item.column()==0:
            try:
                value=int(item.text())
                max_value=get_mkey_bytes[self.param_dict["target_algorithm"]]-1
                if value<0 or value>max_value:
                    raise ValueError(f"请输入正确大小的目标字节！(0-{max_value}")
                for i in range(self.profile_path_model.rowCount()):
                    if item.row()!=i and self.profile_view_bytes[i]==value:
                        raise ValueError("重复的目标字节！")
                self.profile_view_bytes[item.row()]=value
            except Exception as e:
                self.update_text(repr(e),False)
                item.setText(str(self.profile_view_bytes[item.row()]))

    def modify_profile_path(self):
        row = self.profile_path_view.currentIndex().row()
        column = self.profile_path_view.currentIndex().column()
        if column==1:
            fname = QFileDialog.getOpenFileName(self, "Open File", self.param_dict['profile_file_dir'], "*.h5")
            if fname[0]:
                self.profile_path_model.item(row,column).setText(fname[0])

    def get_item_tip(self):
        column = self.profile_path_view.currentIndex().column()
        if column==1:
            path = self.profile_path_view.currentIndex().data()
            QToolTip.showText(QCursor.pos(), path)

    def add_profile_item(self):
        item_tb=None
        row_num = self.profile_path_model.rowCount()
        for i in range(get_mkey_bytes[self.param_dict["target_algorithm"]]):
            if i not in self.profile_view_bytes:
                self.profile_view_bytes.append(i)
                item_tb=QStandardItem(str(i))
                break
        item_tb.setEditable(True)
        item_tb.setTextAlignment(Qt.AlignHCenter|Qt.AlignCenter)
        self.profile_path_model.setItem(row_num,0,item_tb)
        item_path=QStandardItem('未设置')
        item_path.setEditable(False)
        self.profile_path_model.setItem(row_num, 1, item_path)
        if row_num==get_mkey_bytes[self.param_dict["target_algorithm"]]-1:
            self.add_item.setEnabled(False)

    def remove_profile_item(self):
        row = self.profile_path_view.currentIndex().row()
        if row==-1:
            return
        item=self.profile_path_model.item(row,0)
        self.profile_view_bytes.remove(int(item.text()))
        self.profile_path_model.removeRow(row)
        self.add_item.setEnabled(True)

    def update_result_show(self):
        if hasattr(self.attacker,"accuracy") is False:
            return
        show_valid = self.show_valid_result.isChecked()
        val_result_list_len = len(self.attacker.val_accuracy)
        self.figure.clear()
        if show_valid and val_result_list_len > 0:
            ax1 = self.figure.add_subplot(121)
            ax1.set_title('训练集精度与损失')
            ax1.set_ylabel('accuracy')
            ax1.set_xlabel('epochs')
            accuracy,=ax1.plot(self.attacker.accuracy, color='b')
            ax2 = ax1.twinx()
            ax2.set_ylabel('loss')
            loss,=ax2.plot(self.attacker.loss,linestyle='--', color='g')
            ax3 = self.figure.add_subplot(122)
            ax3.set_title('验证集精度与损失')
            ax3.set_ylabel('val_accuracy')
            ax3.set_xlabel('epochs')
            val_accuracy,=ax3.plot(self.attacker.val_accuracy, color='r')
            ax4 = ax3.twinx()
            ax4.set_ylabel('val_loss')
            val_loss,=ax4.plot(self.attacker.val_loss, linestyle='--', color='y')
            self.figure.legend([accuracy,loss],["accuracy","loss"],loc='upper left')
            self.figure.legend([val_accuracy,val_loss],["val_accuracy","val_loss"],loc='upper right')
        else:
            ax1 = self.figure.add_subplot(111)
            ax1.set_title('训练集精度与损失')
            ax1.set_ylabel('accuracy')
            ax1.set_xlabel('epochs')
            accuracy,=ax1.plot(self.attacker.accuracy, color='b')
            ax2 = ax1.twinx()
            ax2.set_ylabel('loss')
            loss,=ax2.plot(self.attacker.loss, linestyle='--', color='g')
            self.figure.legend([accuracy, loss], ["accuracy", "loss"], loc='upper right')
        self.canvas.draw()

    def timer_update_result_show(self):
        cur_result_list_len=len(self.attacker.accuracy)
        if cur_result_list_len==self.result_list_len:
            return
        self.result_list_len=cur_result_list_len
        bar_value=100*cur_result_list_len/self.param_dict['epochs']
        self.update_profile_bar(bar_value)
        self.update_result_show()

    def update_process_text(self,text,finish):
        self.update_text(text,False)
        if finish:
            self.start_profile.setEnabled(True)
            self.attack_method.setEnabled(True)
            self.profile_bar.setValue(0)
            self.profile_bar.hide()
            self.stop_train.setEnabled(False)

    def update_text(self, data, override=True):
        if override:
            self.text_string = ''
        if type(data) is dict:
            self.text_string += 'trs file head:\n'
            for k in data:
                self.text_string += f"  {k}: {data[k]}\n"
            self.text_string += '\n'
        else:
            if data[-1] == '#':
                self.text_string += str(data)[:-1]
            elif len(data)>3 and data[-3:]=='...':
                self.text_string += "(Please Wait)" + str(data) + '\n'
            elif data[-1]=='.':
                self.text_string += "(Execution Success)"+str(data) + '\n'
            elif data[-1]=='!' or data[-1]=='！':
                self.text_string += "(Execution Fail)" + str(data) + '\n'
            else:
                self.text_string += str(data) + '\n'
        self.related_info.setText(self.text_string)
        if len(self.text_string) > 10240:
            self.text_string = self.text_string[7680:]

    def update_profile_bar(self,value):
        self.profile_bar.setValue(value)

    def viewing_dataset_info(self):
        fname = QFileDialog.getOpenFileName(self, "Open file", self.param_dict['profile_file_dir'], "*.h5")
        if fname[0]:
            dataset_file = h5py.File(fname[0], "r")
            description = dataset_file['description'][0].decode('utf-8')
            description = description.replace(',', '\n').replace(',', '\n').replace('{', '{\n').replace('}', '\n}')
            self.update_text(f"数据集信息：\n{description}", False)

    def browse_exist_model_path(self):
        fname = QFileDialog.getOpenFileName(self, "Open file", self.param_dict['model_file_dir'], "*.h5")
        if fname[0]:
            self.continue_train_path=fname[0]
            self.exist_model_path.setText(fname[0])

    def set_train_new_network(self):
        self.continue_train = self.train_new_network.isChecked() is False
        self.set_attack_method(self.param_dict["attack_method"])

    def showEvent(self, a0):
        self.attack_method.setCurrentText(self.param_dict["attack_method"])
        self.continue_train=self.train_new_network.isChecked() is False
        self.set_attack_method(self.param_dict['attack_method'])

    def closeEvent(self, a0):
        self.text_string = ''
        self.update_text("#", True)
        self.startWindow.fresh()
        self.startWindow.show()